import numpy as np
import os

def demonstrate_numpy_io():
    """
    NumPy IO操作演示脚本
    包含菜鸟教程中的所有IO函数示例
    """
    
    print("=== NumPy IO操作演示 ===\n")
    
    # 创建测试数据
    print("1. 创建测试数据")
    a = np.array([1, 2, 3, 4, 5])
    b = np.array([[1, 2, 3], [4, 5, 6]])
    c = np.arange(0, 1.0, 0.1)
    d = np.sin(c)
    
    print("数组a:", a)
    print("数组b:")
    print(b)
    print("数组c:", c)
    print("数组d (sin(c)):", d)
    print("-" * 60)
    
    # 2. numpy.save() - 保存单个数组到.npy文件
    print("2. numpy.save() - 保存单个数组到.npy文件")
    
    # 保存数组a到文件
    np.save('outfile.npy', a)
    print("已将数组a保存到 outfile.npy")
    
    # 不指定扩展名，会自动添加.npy
    np.save('outfile2', a)
    print("已将数组a保存到 outfile2.npy (自动添加扩展名)")
    
    # 检查文件是否存在
    print("检查生成的文件:")
    print("outfile.npy 存在:", os.path.exists('outfile.npy'))
    print("outfile2.npy 存在:", os.path.exists('outfile2.npy'))
    print("-" * 60)
    
    # 3. numpy.load() - 从.npy文件加载数组
    print("3. numpy.load() - 从.npy文件加载数组")
    
    # 加载保存的数组
    b_loaded = np.load('outfile.npy')
    print("从outfile.npy加载的数组:", b_loaded)
    
    # 验证数据是否正确
    print("数据验证 - 原数组与加载数组是否相等:", np.array_equal(a, b_loaded))
    print("-" * 60)
    
    # 4. numpy.savez() - 保存多个数组到.npz文件
    print("4. numpy.savez() - 保存多个数组到.npz文件")
    
    # 保存多个数组，使用关键字参数命名
    np.savez("runoob.npz", b, c, sin_array=d)
    print("已将多个数组保存到 runoob.npz")
    print("数组b作为arr_0，数组c作为arr_1，数组d作为sin_array保存")
    
    # 加载.npz文件
    r = np.load("runoob.npz")
    print("npz文件包含的数组名称:", r.files)
    
    # 访问各个数组
    print("数组b (arr_0):")
    print(r["arr_0"])
    print("数组c (arr_1):", r["arr_1"])
    print("数组d (sin_array):", r["sin_array"])
    
    # 验证数据
    print("数据验证 - 数组b是否相等:", np.array_equal(b, r["arr_0"]))
    print("数据验证 - 数组c是否相等:", np.array_equal(c, r["arr_1"]))
    print("数据验证 - 数组d是否相等:", np.array_equal(d, r["sin_array"]))
    print("-" * 60)
    
    # 5. numpy.savetxt() 和 numpy.loadtxt() - 文本文件读写
    print("5. numpy.savetxt() 和 numpy.loadtxt() - 文本文件读写")
    
    # 5.1 基本文本文件保存和加载
    print("5.1 基本文本文件操作")
    np.savetxt('out.txt', a)
    print("已将数组a保存到 out.txt (默认格式)")
    
    # 加载文本文件
    b_txt = np.loadtxt('out.txt')
    print("从out.txt加载的数组:", b_txt)
    print("数据类型:", b_txt.dtype)
    
    # 5.2 使用分隔符和格式化
    print("\n5.2 使用分隔符和格式化")
    
    # 创建二维数组
    e = np.arange(0, 10, 0.5).reshape(4, -1)
    print("二维数组e:")
    print(e)
    
    # 保存为整数，逗号分隔
    np.savetxt("out_formatted.txt", e, fmt="%d", delimiter=",")
    print("已将数组e保存到 out_formatted.txt (整数格式，逗号分隔)")
    
    # 加载时指定相同的分隔符
    e_loaded = np.loadtxt("out_formatted.txt", delimiter=",")
    print("从out_formatted.txt加载的数组:")
    print(e_loaded)
    print("数据类型:", e_loaded.dtype)
    print("-" * 60)
    
    # 6. 高级文本文件操作示例
    print("6. 高级文本文件操作示例")
    
    # 6.1 自定义格式和头信息
    print("6.1 自定义格式和头信息")
    
    # 创建包含多种数据类型的数组
    data = np.array([
        [1, 2.5, 'hello'],
        [3, 4.8, 'world'],
        [5, 6.2, 'numpy']
    ])
    
    # 使用自定义格式保存
    np.savetxt('data_custom.txt', 
               np.column_stack([np.arange(len(data)), data]), 
               fmt='%d, %.1f, %s', 
               delimiter=' | ',
               header='ID | Value1 | Value2 | Text',
               comments='# ')
    
    print("已创建自定义格式文件: data_custom.txt")
    
    # 6.2 跳过行和选择列
    print("\n6.2 跳过行和选择列")
    
    # 创建带有多行头信息的文件
    with open('data_with_header.txt', 'w') as f:
        f.write("# 这是头信息行1\n")
        f.write("# 这是头信息行2\n")
        f.write("1,2.5,hello\n")
        f.write("3,4.8,world\n")
        f.write("5,6.2,numpy\n")
    
    # 加载时跳过2行头信息
    data_loaded = np.loadtxt('data_with_header.txt', 
                            delimiter=',', 
                            skiprows=2,
                            dtype=str)
    print("跳过2行头信息加载的数据:")
    print(data_loaded)
    print("-" * 60)
    
    # 7. 二进制格式与文本格式对比
    print("7. 二进制格式与文本格式对比")
    
    # 创建大型数组进行对比
    large_array = np.random.rand(1000, 1000)
    
    # 保存为二进制格式
    np.save('large_binary.npy', large_array)
    
    # 保存为文本格式
    np.savetxt('large_text.txt', large_array)
    
    # 检查文件大小
    binary_size = os.path.getsize('large_binary.npy')
    text_size = os.path.getsize('large_text.txt')
    
    print("二进制文件大小: {:.2f} KB".format(binary_size / 1024))
    print("文本文件大小: {:.2f} KB".format(text_size / 1024))
    print("压缩比: {:.2f}%".format((text_size - binary_size) / text_size * 100))
    print("-" * 60)

def cleanup_files():
    """
    清理生成的测试文件
    """
    files_to_remove = [
        'outfile.npy', 'outfile2.npy', 'runoob.npz',
        'out.txt', 'out_formatted.txt', 'data_custom.txt',
        'data_with_header.txt', 'large_binary.npy', 'large_text.txt'
    ]
    
    print("清理测试文件...")
    for file in files_to_remove:
        if os.path.exists(file):
            os.remove(file)
            print(f"已删除: {file}")

if __name__ == "__main__":
    try:
        # 运行演示
        demonstrate_numpy_io()
        
        # 询问是否清理文件
        response = input("\n是否清理生成的测试文件? (y/n): ")
        if response.lower() == 'y':
            cleanup_files()
            print("文件清理完成!")
        else:
            print("测试文件保留在当前目录")
            
    except Exception as e:
        print(f"执行过程中出现错误: {e}")
        print("正在进行文件清理...")
        cleanup_files()
    
    print("\n=== NumPy IO操作演示结束 ===")